import faiss
import numpy as np

#数据
d = 512          #维数
n_data = 2000
np.random.seed(0)
data = []
mu = 3
sigma = 0.1
for i in range(n_data):
    data.append(np.random.normal(mu, sigma, d))
data = np.array(data).astype('float32')

#ids, 6位随机数
ids = []
start = 100000
for i in range(data.shape[0]):
    ids.append(start)
    start += 100
ids = np.array(ids)
# 一般将其作为baseline与其他索引方式对比，以便在精度和时间开销之间做权衡。
# 不支持add_with_ids，如果需要，可以用“IDMap, Flat”。

#不支持add_with_ids
index = faiss.index_factory(d, "Flat")
index.add(data)
dis, ind = index.search(data[:5], 10)
print(ind)

index = faiss.index_factory(d, "IDMap, Flat")
index.add_with_ids(data, ids)
dis, ind = index.search(data[:5], 10)
print(ind)   # 返回的结果是我们自己定义的id


'''
如果稍微有点在意，使用“..., Flat“
"..."是聚类操作，聚类之后将每个向量映射到相应的bucket。该索引类型并不会保存压缩之后的数据，
而是保存原始数据，所以内存开销与原始数据一致。通过nprobe参数控制速度/精度。
支持GPU,但是要注意，选用的聚类操作必须也支持。
faiss.index_factory(d, "IVF100, Flat")
'''


'''
如果很在意，使用”PCARx,...,SQ8“
如果保存全部原始数据的开销太大，可以用这个索引方式。包含三个部分，
1.降维
2.聚类
3.scalar 量化，每个向量编码为8bit 不支持GPU
faiss.index_factory(d, "PCAR16,IVF50,SQ8")
'''

'''
如果非常非常在意，使用"OPQx_y,...,PQx"
y需要是x的倍数，一般保持y<=d，y<=4*x。 支持GPU。
'''